import os
import random
from PIL import Image, ImageDraw, ImageFont
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import gymnasium as gym
from gymnasium import spaces

from gymnasium.envs.registration import register

# 随机选择的中文字集合（可以根据需求进行扩展）
CHINESE_CHARACTERS = list("汉字随机生成测试常用文字")

# 切换到 'TkAgg' 后端，确保图形显示在本地窗口
matplotlib.use('TkAgg')  # 或者使用 'Agg' 后端避免任何显示窗口


class ImageTextGame(gym.Env):
    def __init__(self, image_dir, font_paths, max_text_count=4, render_mode='human'):
        super(ImageTextGame, self).__init__()
        self.image_dir = image_dir
        self.font_paths = font_paths  # 多个字体文件路径的列表
        self.max_text_count = max_text_count  # 限制最大的文字数量
        self.window_width = 260  # 宽度为260
        self.window_height = 160  # 高度保持不变
        self.current_image = None
        self.text_info = []  # 保存多个文字的相关信息
        self.current_step = 0  # 当前的点击步数
        self.total_reward = 0  # 总得分
        self.max_steps = 0  # 最大步数，等于绘制的文字数
        self.click_positions = []  # 记录玩家的点击位置
        self.render_mode = render_mode  # 控制是否可视化
        self.text_string = ""  # 保存绘制的文字串

        # 定义连续的动作空间，坐标范围为 [0, 1]
        self.action_space = spaces.Box(low=np.array([0.0, 0.0]), high=np.array([1.0, 1.0]), dtype=np.float32)

        # 定义观察空间，图像的像素值范围在 0 到 255 之间
        self.observation_space = spaces.Box(low=0, high=255, shape=(self.window_height, self.window_width, 3),
                                            dtype=np.uint8)

    def reset(self, seed=None, options=None):
        # 重置游戏状态
        super().reset(seed=seed)
        self.current_image = self._get_random_image()
        self.text_info = []  # 确保在reset时清空之前的文字信息
        self.current_image = self._draw_random_text(self.current_image)
        self.current_step = 0
        self.total_reward = 0
        self.max_steps = len(self.text_info)  # 最大步数等于绘制的文字数量
        if self.max_steps == 0:
            raise ValueError("max_steps cannot be 0, check if text was drawn correctly.")
        self.click_positions = []  # 清空点击位置记录

        # 返回(图片, 文字字符串)
        return (self.current_image, self.text_string), {}

    def step(self, action):
        '''
        action: [x, y] 坐标形式的动作输入，取值范围为 [0, 1]
        '''
        # print("action:", action)
        x_click_ratio, y_click_ratio = action

        # 将百分比坐标转换为像素坐标
        x_click = x_click_ratio * (self.window_width - 1)
        y_click = y_click_ratio * (self.window_height - 1)

        # 确保点击坐标在图像范围内
        x_click = np.clip(x_click, 0, self.window_width - 1)
        y_click = np.clip(y_click, 0, self.window_height - 1)

        # 检查游戏是否已终止
        if self.current_step >= self.max_steps:
            terminated = True
            truncated = False
            return (self.current_image, self.text_string), self.total_reward, terminated, truncated, {}

        target_info = self.text_info[self.current_step]

        # 获取旋转字符的四个顶点坐标
        rotated_corners = target_info['rotated_corners']

        # 使用射线法判断点击点是否在旋转后的文字区域内
        current_reward = 0
        if self.is_point_in_polygon((x_click, y_click), rotated_corners):
            # 点击正确
            self.total_reward += 2
            current_reward = 2
        else:
            # 点击错误
            self.total_reward -= 2
            current_reward = -2

        # 记录点击位置
        self.click_positions.append((x_click, y_click))

        # 增加步骤
        self.current_step += 1

        # 检查终止条件
        terminated = self.current_step >= self.max_steps
        truncated = False

        if self.render_mode == 'human':
            self.render()

        return (self.current_image, self.text_string), current_reward, terminated, truncated, {}


    def render(self, mode='human'):
        img = Image.fromarray(self.current_image.copy())
        draw = ImageDraw.Draw(img)

        # 标记当前目标字符的旋转边界框
        if self.current_step < len(self.text_info):
            target_info = self.text_info[self.current_step]
            rotated_corners = target_info['rotated_corners']
            draw.polygon([tuple(point) for point in rotated_corners], outline="green")

        # 标记所有点击位置
        for (x_click, y_click) in self.click_positions:
            draw.ellipse([x_click - 5, y_click - 5, x_click + 5, y_click + 5], fill="red", outline="red")

        # 显示图像
        plt.imshow(np.array(img))
        plt.title(f"Step: {self.current_step}, Total Reward: {self.total_reward}")
        plt.axis('off')
        plt.show(block=False)
        plt.pause(0.001)


    def _get_random_image(self):
        # 随机选择目录中的一张图片
        image_files = os.listdir(self.image_dir)
        selected_image_path = os.path.join(self.image_dir, random.choice(image_files))
        image = Image.open(selected_image_path)

        # 调整图片大小
        image = self._adjust_image_size(image)
        return np.array(image)

    def _adjust_image_size(self, image):
        # 获取图片的尺寸
        img_width, img_height = image.size

        if img_width < self.window_width or img_height < self.window_height:
            # 如果图片的大小小于窗口大小，则拉伸图片
            image = image.resize((self.window_width, self.window_height))
        else:
            # 如果图片的大小大于窗口大小，随机截取一个区域
            x_start = random.randint(0, img_width - self.window_width)
            y_start = random.randint(0, img_height - self.window_height)
            image = image.crop((x_start, y_start, x_start + self.window_width, y_start + self.window_height))

        return image

    def _draw_random_text(self, image):
        # 将numpy数组转换为PIL图像
        img_pil = Image.fromarray(image)
        draw = ImageDraw.Draw(img_pil)
        self.text_string = ""  # 重置文本字符串

        # 随机生成指定数量的汉字
        text = random.choices(CHINESE_CHARACTERS, k=self.max_text_count)

        x_position = 0

        for char in text:
            self.text_string += char

            font_size = random.randint(int(self.window_height * 0.15), int(self.window_height * 0.35))
            font_path = random.choice(self.font_paths)
            font = ImageFont.truetype(font_path, font_size)

            color = tuple(np.random.randint(0, 255, size=3))

            # 创建字符图像
            bbox = draw.textbbox((0, 0), char, font=font)
            text_width = bbox[2] - bbox[0]
            text_height = bbox[3] - bbox[1]

            char_img = Image.new("RGBA", (text_width, text_height), (255, 255, 255, 255))
            char_draw = ImageDraw.Draw(char_img)
            char_draw.text((0, 0), char, font=font, fill=color)

            # 旋转字符图像（不裁剪）
            angle = random.randint(-45, 45)
            rotated_char_img = char_img.rotate(angle, expand=1)

            # 获取旋转后图像的尺寸
            rotated_width, rotated_height = rotated_char_img.size

            # 调整 x_position 和 y_position 以确保字符在图像边界内
            if x_position + rotated_width > self.window_width:
                x_position = self.window_width - rotated_width
            if x_position < 0:
                x_position = 0

            y_position = random.randint(0, self.window_height - rotated_height)
            if y_position + rotated_height > self.window_height:
                y_position = self.window_height - rotated_height
            if y_position < 0:
                y_position = 0

            # 将旋转后的字符图像粘贴到主图像上
            img_pil.paste(rotated_char_img, (x_position, y_position), rotated_char_img)

            # 计算旋转后的四个顶点坐标
            w, h = text_width, text_height
            # 原始字符图像的四个角点
            corners = np.array([
                [0, 0],
                [w, 0],
                [w, h],
                [0, h]
            ])

            # 计算旋转中心（字符图像的中心）
            center = np.array([w / 2, h / 2])

            # 构造旋转矩阵
            theta = np.radians(angle)
            rotation_matrix = np.array([
                [np.cos(theta), -np.sin(theta)],
                [np.sin(theta), np.cos(theta)]
            ])

            # 旋转后的角点坐标
            rotated_corners = np.dot(corners - center, rotation_matrix) + center

            # 考虑旋转后图像的尺寸变化，计算偏移量
            offset_x = (rotated_width - w) / 2
            offset_y = (rotated_height - h) / 2

            # 将角点坐标平移到旋转后图像的坐标系中
            rotated_corners += np.array([offset_x, offset_y])

            # 最终的角点坐标，平移到主图像上的位置
            rotated_corners += np.array([x_position, y_position])

            # 记录字符信息，包括旋转后的顶点坐标
            self.text_info.append({
                "char": char,
                "font_size": font_size,
                "color": color,
                "position": (x_position, y_position),
                "font_path": font_path,
                "rotation_angle": angle,
                "width": rotated_width,
                "height": rotated_height,
                "rotated_corners": rotated_corners  # 添加旋转后的角点坐标
            })

            x_position += rotated_width + 10  # 增加一些间距

        return np.array(img_pil)


    def get_target_area(self):
        if self.current_step < len(self.text_info):
            target_info = self.text_info[self.current_step]

            rotated_corners = target_info['rotated_corners']

            # 获取旋转字符区域的最小和最大坐标
            x_min = np.min(rotated_corners[:, 0])
            x_max = np.max(rotated_corners[:, 0])
            y_min = np.min(rotated_corners[:, 1])
            y_max = np.max(rotated_corners[:, 1])

            # 不断生成随机点，直到点在旋转矩形内
            max_attempts = 10000
            for _ in range(max_attempts):
                x_click = random.uniform(x_min, x_max)
                y_click = random.uniform(y_min, y_max)
                if self.is_point_in_polygon((x_click, y_click), rotated_corners):
                    # 将像素坐标转换为百分比坐标
                    x_click_ratio = x_click / (self.window_width - 1)
                    y_click_ratio = y_click / (self.window_height - 1)
                    return (self.current_image, self.text_string), np.array([x_click_ratio, y_click_ratio])
            
            raise ValueError("Failed to generate target area within the polygon.")

    def is_point_in_polygon(self, point, polygon):
        # 获取待检测点的坐标
        x, y = point
        # 获取多边形顶点的数量
        num = len(polygon)
        # 初始化内部标志，默认在外边
        inside = False
        # 获取多边形的第一个顶点
        p1x, p1y = polygon[0]
        # 遍历多边形的边
        for i in range(num + 1):
            # 获取当前边的结束顶点，第一个顶点可能是和polygon[0]相同的点
            p2x, p2y = polygon[i % num]
            # 检查Y是否在当前Y范围内
            if min(p1y, p2y) < y <= max(p1y, p2y):
                # 检查X是否在当前边最大X的右侧
                if x <= max(p1x, p2x):
                    # 当前边不是水平的情况下
                    if p1y != p2y:
                        # 计算当前射线和边的交点的X坐标 todo 了解数学原理
                        xinters = (y - p1y) * (p2x - p1x) / (p2y - p1y) + p1x
                    if p1x == p2x or x <= xinters:
                        inside = not inside
            p1x, p1y = p2x, p2y
        return inside



def save_obs(obs, size=None):
    # 保存obs中的图片，图片名字以obs中的文字串命名
    image, text_string = obs
    img = Image.fromarray(image)

    if size:
        img = img.resize(size, Image.LANCZOS)

    filename = f"{text_string}.png"
    img.save(filename)


def test_get_target_area():
    env = ImageTextGame(
        image_dir=r"D:\Projects\DataSet\gym\wordVision",
        font_paths=[r"D:\Projects\DataSet\gym\font\simsun.ttc"],
        render_mode="human"
    )
    for _ in range(1000):
        obs, _ = env.reset()
        for i in range(env.max_steps):
            _, action = env.get_target_area()  # 获取目标区域内的点击坐标（百分比形式）
            obs, reward, _, _, _ = env.step(action)
            # 保存obs中的图片，图片名字以obs中的文字串命名
            save_obs(obs)
            print(f"Current reward: {reward}")
            if reward < 0:
                raise ValueError("Game Over")

    print("Test Pass")


def test_env():
    env = ImageTextGame(
        image_dir=r"D:\Projects\datasets\gym\wordVision",
        font_paths=[r"D:\Projects\datasets\gym\font\simsun.ttc"],
        render_mode="human"
    )
    for _ in range(1):
        obs, _ = env.reset()
        for i in range(env.max_steps):
            action = env.action_space.sample()  # 随机生成点击位置（百分比形式）
            obs, reward, terminated, truncated, _ = env.step(action)
            save_obs(obs, size=(84, 84))
            print(f"Current reward: {reward}")
            if terminated or truncated:
                print("Game Over")
                break

    print("Text Info:", env.text_info)  # 输出绘制的文字信息


# 用法示例
if __name__ == "__main__":
    # 注册自定义环境，id是唯一标识
    register(
        id='ImageTextGame-v0',  # 环境ID
        entry_point='game.ImageTextGame',  # 指向你环境的入口，即类的导入路径
        max_episode_steps=4,  # 每一回合最大步数
        kwargs={
            'image_dir': r"D:\Projects\datasets\gym\wordVision",  # 初始化时传入的参数
            'font_paths': [r"D:\Projects\datasets\gym\font\simsun.ttc"],
            'max_text_count': 4,
            'render_mode': 'service'
        }
    )

    # test_env()
    test_get_target_area()
